'''
基于内存的对话历史管理 - 轻量级会话存储方案

内存存储的特点：
✅ 优点：
   - 零依赖，无需外部数据库
   - 读写速度极快（内存操作）
   - 实现简单，适合开发和测试
   - 无网络延迟

❌ 缺点：
   - 服务重启后数据丢失
   - 单机部署，不支持分布式
   - 内存占用随对话增多而增加
   - 不适合生产环境大量用户

适用场景：
- 本地开发和测试
- 原型验证和演示
- 单用户桌面应用
- 短期会话场景
'''
import os

from dotenv import load_dotenv
from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_openai import ChatOpenAI

# 自动加载 .env 文件中的环境变量
load_dotenv()

# ========== 构建基础对话处理链 ==========
'''
创建聊天提示词模板，定义对话流程：
1. system: 设定AI的角色和能力范围
2. history: 历史消息占位符（从内存存储加载）
3. human: 当前用户输入
'''
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个优秀的助手，你的能力是{ability}，每次你返回问题的结果不要超过30个字"),
        MessagesPlaceholder(variable_name="history"),  # 历史消息注入点
        ("human", "{input}")  # 当前用户输入
    ]
)

# 初始化大语言模型
llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen3-32B", streaming=True, temperature=0)

# 输出解析器：将模型输出转换为纯文本字符串
parser = StrOutputParser()

# 构建完整处理链：提示词 -> 大模型 -> 输出解析
runnable = prompt | llm | parser

# ========== 内存存储管理 ==========
'''
内存存储设计：
使用Python字典作为存储后端，以session_id为键，InMemoryChatMessageHistory为值

存储结构示例：
{
    "lqq_1": InMemoryChatMessageHistory([message1, message2, ...]),
    "lqq_2": InMemoryChatMessageHistory([message3, message4, ...]),
    "user_123": InMemoryChatMessageHistory([message5, message6, ...])
}
'''
store = {}  # 全局内存存储字典


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    """
    根据session_id获取或创建对话历史

    工作流程：
    1. 检查session_id是否已存在于store中
    2. 如果存在，返回对应的历史记录
    3. 如果不存在，创建新的InMemoryChatMessageHistory并存储

    参数：
        session_id: 会话唯一标识符，用于隔离不同对话

    返回：
        BaseChatMessageHistory实例，用于操作该会话的历史消息
    """
    # 如果session_id不存在于存储中，创建新的内存历史记录
    if session_id not in store:
        print(f"🆕 创建新的会话历史: {session_id}")
        store[session_id] = InMemoryChatMessageHistory()
    else:
        print(f"📚 使用现有会话历史: {session_id}")
        print(f"   当前消息数: {len(store[session_id].messages)}")

    return store[session_id]


# ========== 创建带历史管理的执行器 ==========
'''
RunnableWithMessageHistory 自动管理对话历史：
- 根据config中的session_id获取对应历史
- 自动将历史消息注入到提示词中
- 自动保存新的对话记录到内存存储
'''
with_message_history = RunnableWithMessageHistory(
    runnable,  # 基础处理链
    get_session_history,  # 历史记录获取函数
    input_messages_key="input",  # 用户输入的参数名
    history_messages_key="history"  # 历史消息在提示词中的参数名
)

# ========== 演示内存存储的实际效果 ==========
print("🚀 内存存储对话历史演示")
print("=" * 50)

# 第一次对话：创建新会话 lqq_1
print("\n💬 第一次对话 - 创建新会话 lqq_1")
response1 = with_message_history.invoke(
    {"ability": "历史", "input": "清朝第一个皇帝是谁？"},
    config={"configurable": {"session_id": "lqq_1"}}
)
print(f"🤖 AI回复: {response1}")

# 显示当前存储状态
print(f"\n📊 存储状态: 共有 {len(store)} 个会话")
for session_id, history in store.items():
    print(f"   📁 {session_id}: {len(history.messages)} 条消息")

# 第二次对话：使用相同会话 lqq_1（应该有历史记忆）
print("\n💬 第二次对话 - 使用相同会话 lqq_1（测试历史记忆）")
response2 = with_message_history.invoke(
    {"ability": "历史", "input": "什么？"},  # 模糊问题，依赖历史上下文
    config={"configurable": {"session_id": "lqq_1"}}
)
print(f"🤖 AI回复: {response2}")

# 显示历史消息内容
print(f"\n📖 会话 lqq_1 的完整历史:")
for i, message in enumerate(store["lqq_1"].messages):
    print(f"   {i + 1}. {message.type}: {message.content}")

# 第三次对话：使用新会话 lqq_2（应该没有历史记忆）
print("\n💬 第三次对话 - 创建新会话 lqq_2（测试会话隔离）")
response3 = with_message_history.invoke(
    {"ability": "历史", "input": "这个朝代有几个皇帝？"},  # 同样的问题，但新会话
    config={"configurable": {"session_id": "lqq_2"}}
)
print(f"🤖 AI回复: {response3}")

# 最终存储状态
print(f"\n🎯 最终存储状态:")
print(f"   总会话数: {len(store)}")
for session_id, history in store.items():
    print(f"   {session_id}: {len(history.messages)} 条消息")


# ========== 进阶功能：存储分析和工具函数 ==========
def analyze_storage():
    """分析内存存储的使用情况"""
    print("\n🔍 内存存储分析")
    print("=" * 30)

    total_messages = 0
    for session_id, history in store.items():
        msg_count = len(history.messages)
        total_messages += msg_count
        print(f"   {session_id}: {msg_count} 条消息")

    print(f"\n📈 统计摘要:")
    print(f"   会话数量: {len(store)}")
    print(f"   总消息数: {total_messages}")
    print(f"   平均消息/会话: {total_messages / len(store) if store else 0:.1f}")


def clear_session(session_id: str):
    """清空指定会话的历史"""
    if session_id in store:
        store[session_id].clear()
        print(f"✅ 已清空会话 {session_id} 的历史")
    else:
        print(f"❌ 会话 {session_id} 不存在")


def export_session_history(session_id: str):
    """导出会话历史为可读格式"""
    if session_id in store:
        history = store[session_id]
        print(f"\n📤 会话 {session_id} 的历史导出:")
        for i, message in enumerate(history.messages):
            role = "👤 用户" if message.type == "human" else "🤖 AI"
            print(f"   {i + 1}. {role}: {message.content}")
    else:
        print(f"❌ 会话 {session_id} 不存在")


# 运行分析工具
analyze_storage()

# 导出会话历史查看
export_session_history("lqq_1")

# ========== 内存存储 vs Redis存储 对比 ==========
'''
📊 内存存储 vs Redis存储 对比表：

| 特性         | 内存存储                    | Redis存储                    |
|--------------|-----------------------------|------------------------------|
| 数据持久性   | ❌ 服务重启丢失              | ✅ 持久化到磁盘               |
| 部署复杂度   | ✅ 零依赖，简单              | ❌ 需要Redis服务              |
| 性能         | ✅ 内存级速度               | ✅ 内存级速度                 |
| 分布式支持   | ❌ 单机 only                | ✅ 支持多实例                 |
| 适用场景     | 开发、测试、演示            | 生产环境、多用户系统         |
| 数据隔离     | ✅ 进程内隔离               | ✅ 数据库级隔离               |
| 内存管理     | ❌ 需手动清理               | ✅ 自动TTL过期               |

选择建议：
- 开发测试 → 内存存储（简单快捷）
- 生产环境 → Redis存储（稳定可靠）
- 单机应用 → 内存存储（轻量高效）  
- 分布式系统 → Redis存储（共享存储）
'''


# ========== 实际应用场景示例 ==========
def demonstrate_use_cases():
    """演示不同场景下的内存存储应用"""

    print("\n🎯 内存存储实际应用场景")
    print("=" * 40)

    # 场景1：临时客服对话
    print("\n1. 🏢 临时客服对话")
    store["customer_001"] = InMemoryChatMessageHistory()
    response = with_message_history.invoke(
        {"ability": "客服", "input": "我的订单有问题"},
        config={"configurable": {"session_id": "customer_001"}}
    )
    print(f"   客服回复: {response}")

    # 场景2：个人学习助手
    print("\n2. 📚 个人学习助手")
    store["study_session"] = InMemoryChatMessageHistory()
    response = with_message_history.invoke(
        {"ability": "数学辅导", "input": "解释勾股定理"},
        config={"configurable": {"session_id": "study_session"}}
    )
    print(f"   学习助手: {response}")

    # 场景3：多轮对话测试
    print("\n3. 🔄 多轮对话连续性测试")
    test_session = "multi_turn_test"
    questions = [
        "我叫张三",
        "记住我的名字了吗？",
        "我的名字是什么？"
    ]

    for i, question in enumerate(questions, 1):
        response = with_message_history.invoke(
            {"ability": "记忆测试", "input": question},
            config={"configurable": {"session_id": test_session}}
        )
        print(f"   第{i}轮: {response}")


# 取消注释运行场景演示
# demonstrate_use_cases()

print("\n" + "=" * 50)
print("🎉 内存存储演示完成!")
print("=" * 50)
print("💡 核心要点总结:")
print("   1. session_id 是会话隔离的关键")
print("   2. 相同session_id保持对话连续性")
print("   3. 不同session_id完全独立隔离")
print("   4. 内存存储适合开发和测试环境")
print("   5. 生产环境建议使用Redis等持久化存储")